summary network
From Mice to Trains: Amortized Bayesian Inference on Graph Data
Jedhoff, Svenja, Semenova, Elizaveta, Raulo, Aura, Meyer, Anne, Bürkner, Paul-Christian
Graphs arise across diverse domains, from biology and chemistry to social and information networks, as well as in transportation and logistics. Inference on graph-structured data requires methods that are permutation-invariant, scalable across varying sizes and sparsities, and capable of capturing complex long-range dependencies, making posterior estimation on graph parameters particularly challenging. Amortized Bayesian Inference (ABI) is a simulation-based framework that employs generative neural networks to enable fast, likelihood-free posterior inference. We adapt ABI to graph data to address these challenges to perform inference on node-, edge-, and graph-level parameters. Our approach couples permutation-invariant graph encoders with flexible neural posterior estimators in a two-module pipeline: a summary network maps attributed graphs to fixed-length representations, and an inference network approximates the posterior over parameters. In this setting, several neural architectures can serve as the summary network. In this work we evaluate multiple architectures and assess their performance on controlled synthetic settings and two real-world domains -- biology and logistics -- in terms of recovery and calibration.
Robust Amortized Bayesian Inference with Self-Consistency Losses on Unlabeled Data
Mishra, Aayush, Habermann, Daniel, Schmitt, Marvin, Radev, Stefan T., Bürkner, Paul-Christian
Neural amortized Bayesian inference (ABI) can solve probabilistic inverse problems orders of magnitude faster than classical methods. However, neural ABI is not yet sufficiently robust for widespread and safe applicability. In particular, when performing inference on observations outside of the scope of the simulated data seen during training, for example, because of model misspecification, the posterior approximations are likely to become highly biased. Due to the bad pre-asymptotic behavior of current neural posterior estimators in the out-of-simulation regime, the resulting estimation biases cannot be fixed in acceptable time by just simulating more training data. In this proof-of-concept paper, we propose a semi-supervised approach that enables training not only on (labeled) simulated data generated from the model, but also on unlabeled data originating from any source, including real-world data. To achieve the latter, we exploit Bayesian self-consistency properties that can be transformed into strictly proper losses without requiring knowledge of true parameter values, that is, without requiring data labels. The results of our initial experiments show remarkable improvements in the robustness of ABI on out-of-simulation data. Even if the observed data is far away from both labeled and unlabeled training data, inference remains highly accurate. If our findings also generalize to other scenarios and model classes, we believe that our new method represents a major breakthrough in neural ABI.
Amortized Bayesian Mixture Models
Kucharský, Šimon, Bürkner, Paul Christian
Finite mixtures are a broad class of models useful in scenarios where observed data is generated by multiple distinct processes but without explicit information about the responsible process for each data point. Estimating Bayesian mixture models is computationally challenging due to issues such as high-dimensional posterior inference and label switching. Furthermore, traditional methods such as MCMC are applicable only if the likelihoods for each mixture component are analytically tractable. Amortized Bayesian Inference (ABI) is a simulation-based framework for estimating Bayesian models using generative neural networks. This allows the fitting of models without explicit likelihoods, and provides fast inference. ABI is therefore an attractive framework for estimating mixture models. This paper introduces a novel extension of ABI tailored to mixture models. We factorize the posterior into a distribution of the parameters and a distribution of (categorical) mixture indicators, which allows us to use a combination of generative neural networks for parameter inference, and classification networks for mixture membership identification. The proposed framework accommodates both independent and dependent mixture models, enabling filtering and smoothing. We validate and demonstrate our approach through synthetic and real-world datasets.
Solving High-dimensional Inverse Problems Using Amortized Likelihood-free Inference with Noisy and Incomplete Data
Zeng, Jice, Wang, Yuanzhe, Tartakovsky, Alexandre M., Barajas-Solano, David
We present a likelihood-free probabilistic inversion method based on normalizing flows for high-dimensional inverse problems. The proposed method is composed of two complementary networks: a summary network for data compression and an inference network for parameter estimation. The summary network encodes raw observations into a fixed-size vector of summary features, while the inference network generates samples of the approximate posterior distribution of the model parameters based on these summary features. The posterior samples are produced in a deep generative fashion by sampling from a latent Gaussian distribution and passing these samples through an invertible transformation. We construct this invertible transformation by sequentially alternating conditional invertible neural network and conditional neural spline flow layers. The summary and inference networks are trained simultaneously. We apply the proposed method to an inversion problem in groundwater hydrology to estimate the posterior distribution of the log-conductivity field conditioned on spatially sparse time-series observations of the system's hydraulic head responses.The conductivity field is represented with 706 degrees of freedom in the considered problem.The comparison with the likelihood-based iterative ensemble smoother PEST-IES method demonstrates that the proposed method accurately estimates the parameter posterior distribution and the observations' predictive posterior distribution at a fraction of the inference time of PEST-IES.
Conditional diffusions for neural posterior estimation
Chen, Tianyu, Bansal, Vansh, Scott, James G.
Neural posterior estimation (NPE), a simulation-based computational approach for Bayesian inference, has shown great success in situations where posteriors are intractable or likelihood functions are treated as "black boxes." Existing NPE methods typically rely on normalizing flows, which transform a base distributions into a complex posterior by composing many simple, invertible transformations. But flow-based models, while state of the art for NPE, are known to suffer from several limitations, including training instability and sharp trade-offs between representational power and computational cost. In this work, we demonstrate the effectiveness of conditional diffusions as an alternative to normalizing flows for NPE. Conditional diffusions address many of the challenges faced by flow-based methods. Our results show that, across a highly varied suite of benchmarking problems for NPE architectures, diffusions offer improved stability, superior accuracy, and faster training times, even with simpler, shallower models. These gains persist across a variety of different encoder or "summary network" architectures, as well as in situations where no summary network is required.
Amortized Bayesian Multilevel Models
Habermann, Daniel, Schmitt, Marvin, Kühmichel, Lars, Bulling, Andreas, Radev, Stefan T., Bürkner, Paul-Christian
Obtaining accurate inference and faithful uncertainty quantification in reasonable time is a frontier of today's statistical research (Cranmer et al., 2020). One major difficulty arising in most experimental and almost all observational data is the presence of complex dependency structures, for example, due to natural groupings (e.g., data gathered in different countries) or repeated measurements of the same observational units over time (e.g., particles, bacteria, or people; Gelman and Hill, 2006). To leverage these dependency structures, multilevel models (MLMs), also referred to as latent variable, hierarchical, random, or mixed effects models, have become an integral part of modern Bayesian statistics (Goldstein, 2011; Gelman et al., 2013; McGlothlin and Viele, 2018; Finch et al., 2019; Yao et al., 2022). Despite the wide success of Bayesian MLMs across the quantitative sciences, a major challenge is their limited efficiency and scalability when dealing with large and complex data. This is because estimating the full posterior distribution of all parameters of interest can be very costly (Gelman et al., 2013).
Detecting Model Misspecification in Amortized Bayesian Inference with Neural Networks: An Extended Investigation
Schmitt, Marvin, Bürkner, Paul-Christian, Köthe, Ullrich, Radev, Stefan T.
Recent advances in probabilistic deep learning enable efficient amortized Bayesian inference in settings where the likelihood function is only implicitly defined by a simulation program (simulation-based inference; SBI). But how faithful is such inference if the simulation represents reality somewhat inaccurately, that is, if the true system behavior at test time deviates from the one seen during training? We conceptualize the types of such model misspecification arising in SBI and systematically investigate how the performance of neural posterior approximators gradually deteriorates as a consequence, making inference results less and less trustworthy. To notify users about this problem, we propose a new misspecification measure that can be trained in an unsupervised fashion (i.e., without training data from the true distribution) and reliably detects model misspecification at test time. Our experiments clearly demonstrate the utility of our new measure both on toy examples with an analytical ground-truth and on representative scientific tasks in cell biology, cognitive decision making, disease outbreak dynamics, and computer vision. We show how the proposed misspecification test warns users about suspicious outputs, raises an alarm when predictions are not trustworthy, and guides model designers in their search for better simulators.
Variational Bayesian Optimal Experimental Design with Normalizing Flows
Dong, Jiayuan, Jacobsen, Christian, Khalloufi, Mehdi, Akram, Maryam, Liu, Wanjiao, Duraisamy, Karthik, Huan, Xun
Bayesian optimal experimental design (OED) seeks experiments that maximize the expected information gain (EIG) in model parameters. Directly estimating the EIG using nested Monte Carlo is computationally expensive and requires an explicit likelihood. Variational OED (vOED), in contrast, estimates a lower bound of the EIG without likelihood evaluations by approximating the posterior distributions with variational forms, and then tightens the bound by optimizing its variational parameters. We introduce the use of normalizing flows (NFs) for representing variational distributions in vOED; we call this approach vOED-NFs. Specifically, we adopt NFs with a conditional invertible neural network architecture built from compositions of coupling layers, and enhanced with a summary network for data dimension reduction. We present Monte Carlo estimators to the lower bound along with gradient expressions to enable a gradient-based simultaneous optimization of the variational parameters and the design variables. The vOED-NFs algorithm is then validated in two benchmark problems, and demonstrated on a partial differential equation-governed application of cathodic electrophoretic deposition and an implicit likelihood case with stochastic modeling of aphid population. The findings suggest that a composition of 4--5 coupling layers is able to achieve lower EIG estimation bias, under a fixed budget of forward model runs, compared to previous approaches. The resulting NFs produce approximate posteriors that agree well with the true posteriors, able to capture non-Gaussian and multi-modal features effectively.
BayesFlow: Amortized Bayesian Workflows With Neural Networks
Radev, Stefan T, Schmitt, Marvin, Schumacher, Lukas, Elsemüller, Lasse, Pratz, Valentin, Schälte, Yannik, Köthe, Ullrich, Bürkner, Paul-Christian
Modern Bayesian inference involves a mixture of computational techniques for estimating, validating, and drawing conclusions from probabilistic models as part of principled workflows for data analysis (Bürkner et al., 2022; Gelman et al., 2020; Schad et al., 2021). Typical problems in Bayesian workflows are the approximation of intractable posterior distributions for diverse model types and the comparison of competing models of the same process in terms of their complexity and predictive performance. However, despite their theoretical appeal and utility, the practical execution of Bayesian workflows is often limited by computational bottlenecks: Obtaining even a single posterior may already take a long time, such that repeated estimation for the purpose of model validation or calibration becomes completely infeasible. BayesFlow provides a framework for simulation-based training of established neural network architectures, such as transformers (Vaswani et al., 2017) and normalizing flows (Papamakarios et al., 2021), for amortized data compression and inference. Amortized Bayesian inference (ABI), as implemented in BayesFlow, enables users to train custom neural networks on model simulations and re-use these networks for any subsequent application of the models. Since the trained networks can perform inference almost instantaneously (typically well below one second), the upfront neural network training is quickly amortized. For instance, amortized inference allows us to test a model's ability to recover its parameters (Schad et al., 2021) or assess its simulation-based calibration (Säilynoja et al., 2022; Talts et al., 2018) for different data set sizes in a matter of seconds, even though this may require the estimation of thousands of posterior distributions. BayesFlow offers a user-friendly API, which encapsulates the details of neural network architectures and training procedures that are less relevant for the practitioner and provides robust default implementations that work well across many applications. At the same time, BayesFlow implements a modular software architecture, allowing machine learning scientists to modify every component of the pipeline for custom applications as well as research at the frontier of Bayesian inference.
Variational inference of fractional Brownian motion with linear computational complexity
Verdier, Hippolyte, Laurent, François, Cassé, Alhassan, Vestergaard, Christian, Masson, Jean-Baptiste
We introduce a simulation-based, amortised Bayesian inference scheme to infer the parameters of random walks. Our approach learns the posterior distribution of the walks' parameters with a likelihood-free method. In the first step a graph neural network is trained on simulated data to learn optimized low-dimensional summary statistics of the random walk. In the second step an invertible neural network generates the posterior distribution of the parameters from the learnt summary statistics using variational inference. We apply our method to infer the parameters of the fractional Brownian motion model from single trajectories. The computational complexity of the amortized inference procedure scales linearly with trajectory length, and its precision scales similarly to the Cram{\'e}r-Rao bound over a wide range of lengths. The approach is robust to positional noise, and generalizes well to trajectories longer than those seen during training. Finally, we adapt this scheme to show that a finite decorrelation time in the environment can furthermore be inferred from individual trajectories.